import os
import dgl
import time
import random
import logging
import itertools
import numpy as np
import pandas as pd

import torch
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

from Config import Config, get_model_info
from nets.load_net import gnn_model
from data.data_set import mRNADataset, collate_fn
from train.train_mRNA import train_network, evaluate_network

"""
GPU设置
"""


def gpu_setup(use_gpu, gpu_id):
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:', torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device


"""
定义log文件记录数据
"""


def get_logger(log_file):
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )
    logger = logging.getLogger()
    return logger


"""
训练验证测试的pipeline
"""


def train_val_pipeline(model_name, config):
    # 记录总的训练时间
    t0 = time.time()
    # 每个epoch的时间列表
    per_epoch_time = []
    """GPU设置"""
    device = config.device

    """定义log文件保存数据"""
    logger = None
    if config.need_log:
        logger = get_logger(config.log_file)

    """设置随机种子"""
    # setting seeds 设置随机种子，保证每次训练随机参数一致
    random.seed(config.seed)
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    if device.type == 'cuda':
        torch.cuda.manual_seed(config.seed)

    """定义模型"""
    # 定义模型
    model = gnn_model(model_name, config)
    model = model.to(device)
    # 显示模型参数
    if config.need_log:
        logger.info(get_model_info(model_name, config))  # 使用 log 输出
    else:
        print(get_model_info(model_name, config))

    # 选用Adam优化器
    optimizer = optim.Adam(model.parameters(), lr=config.init_lr, weight_decay=config.weight_decay)

    scheduler = None
    if config.lr_decay:
        # 当参考的评价指标停止改进时,降低学习率,factor为每次下降的比例,
        # 训练过程中,当指标连续patience次数还没有改进时,降低学习率
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
                                                         factor=config.lr_reduce_factor,
                                                         patience=config.lr_schedule_patience,
                                                         verbose=True)

    """读取数据"""
    # 读取训练/验证/测试数据
    config = Config(model_name, device)
    train_set = mRNADataset(config.train_file, config, mode='train')
    val_set = mRNADataset(config.val_file, config, mode='val')
    test_set = mRNADataset(config.test_file, config, mode='test')

    if config.pos_enc:  # 是否使用位置编码
        t1 = time.time()
        train_set._add_positional_encodings(pos_enc_dim=config.pos_enc_dim)
        val_set._add_positional_encodings(pos_enc_dim=config.pos_enc_dim)
        infor = "已使用位置编码，本次编码耗时:{:.3f} s\n".format(time.time() - t1)
        if config.need_log:
            logger.info(infor)  # 使用 log 输出
        else:
            print(infor)

    # 使用data_loader读取数据
    train_loader = DataLoader(train_set, batch_size=config.batch_size, shuffle=True, collate_fn=collate_fn())
    val_loader = DataLoader(val_set, batch_size=config.batch_size, shuffle=False, collate_fn=collate_fn())
    test_loader = DataLoader(test_set, batch_size=config.batch_size, shuffle=False, collate_fn=collate_fn())

    """训练过程"""
    # 在训练过程中可使用 Ctrl + C 跳出训练
    try:

        for epoch in range(config.num_epoch):
            # 记录开始时间
            start = time.time()

            # 训练函数，返回loss值和评价指标
            epoch_train_loss, epoch_train_assess, optimizer = train_network(model, optimizer, device, train_loader,
                                                                            config)
            # 获取验证loss和评价指标
            epoch_val_loss, epoch_val_assess = evaluate_network(model, device, val_loader, config)
            # 获取测试评价指标
            # _, epoch_test_assess = evaluate_network(model, device, test_loader, config)

            if config.lr_decay:
                # 根据epoch_val_loss降低学习率
                scheduler.step(epoch_val_loss)

            # 保存每轮epoch的耗费时间
            epoch_time = time.time() - start
            per_epoch_time.append(epoch_time)

            # 打印本轮epoch的信息
            infor = "\n" + '-' * 42 + "\n"
            infor += "epoch:{:d}\tepoch_time:{:.5f} s\n".format(epoch + 1, epoch_time)
            infor += "train_loss:{:.5f}, train_mcrmse:{:.5f}\n".format(epoch_train_loss, epoch_train_assess)
            infor += "val_loss:  {:.5f}, val_mcrmse:  {:.5f}\n".format(epoch_val_loss, epoch_val_assess)
            # 由于test集没有ground truth 所以不再显示
            # infor += ">>>>>>>>>>>>>>>>>>> test_mcrmse: {:.5f}\n".format(epoch_test_assess)
            infor += '-' * 42
            if config.need_log:
                logger.info(infor)  # 使用 log 输出
            else:
                print(infor)
            # 是否保存checkpoint的权重
            if config.save_model:
                torch.save(model.state_dict(), '{}.pkl'.format(config.model_path + "/epoch_" + str(epoch + 1)))

    except KeyboardInterrupt:
        infor = "\n" + '-' * 42 + "\n"
        infor += "已使用 Ctrl + C 将训练提前退出"
        if config.need_log:
            logger.info(infor)  # 使用 log 输出
        else:
            print(infor)

    # 训练结束后打印训练情况
    _, val_assess = evaluate_network(model, device, val_loader, config)
    _, train_assess = evaluate_network(model, device, train_loader, config)
    # log输出
    infor = "\n训练 MCR-MSE:\t{:.5f}\n".format(train_assess)
    infor += "验证 MCR-MSE:\t{:.5f}\n".format(val_assess)
    infor += "训练总用时:\t{:.5f} s\n".format(time.time() - t0)
    infor += "平均每个epoch用时:\t{:.5f} s\n".format(np.mean(per_epoch_time))
    if config.need_log:
        logger.info(infor)  # 使用 log 输出
    else:
        print(infor)


def main():
    # 使用模型类型
    # model_name = 'GCN'
    # model_name = 'GAT'
    model_name = 'GatedGCN'
    device = gpu_setup(True, 0)  # 使用设备(cpu/gpu)
    config = Config(model_name, device)  # 参数设置
    train_val_pipeline(model_name, config)  # 进行训练


if __name__ == '__main__':
    main()
